# -*- coding: utf-8 -*-
# @Author : HuangGaoShuang
# @Time : 2024/6/20 9:24

import os
from glob import glob
import numpy as np
import pickle
from tqdm import tqdm

def calculate_center_distance(center_gt, centers_pred):
    """
    计算一个中心点与多个中心点之间的欧几里得距离。
    """
    center_gt = center_gt.reshape(1, -1)  # 确保gt_center为二维数组
    distance = np.sqrt(np.sum((centers_pred - center_gt)**2, axis=1))
    return distance


def load_reference_images(reference_path):
    """
    加载所有参考图像的路径及其坐标信息。
    参数:
        reference_path (str): 参考图像所在的目录路径。
    返回:
        dict: 以索引为键，值为包含路径和坐标的字典列表。
    """
    references = {}
    for path in glob(os.path.join(reference_path, '*')):
        try:
            filename = os.path.basename(path)
            index, x1, y1, x2, y2 = filename[:-4].split('_')
            references.setdefault(index, []).append({
                'name' : filename,
                'path': path,
                'coords': np.array([int(x1), int(y1), int(x2), int(y2)])
            })
        except ValueError:
            print(f"Warning: Skipping file {path}, filename format is incorrect.")
    return references


def process_query_image_center(query_image_name, reference_images, topk=128):
    """
    处理单个查询图像，计算其与参考图像的中心距离，并返回排序后的距离列表。
    """
    index, x1, y1, x2, y2 = query_image_name[:-4].split('_')
    # 将字符串转换为整数并计算中心点坐标
    center_x = (int(x1) + int(x2)) / 2
    center_y = (int(y1) + int(y2)) / 2
    gt_center = np.array([[center_x, center_y]])

    if index in reference_images:
        pred_boxes = np.array([ref['coords'] for ref in reference_images[index]])
        # 分别计算x和y坐标的平均值
        pred_centers_x = (pred_boxes[:, 0] + pred_boxes[:, 2]) / 2
        pred_centers_y = (pred_boxes[:, 1] + pred_boxes[:, 3]) / 2
        pred_centers = np.column_stack((pred_centers_x, pred_centers_y))  # 创建中心点坐标数组
        distances = calculate_center_distance(gt_center, pred_centers)

        # 获取前top_k个最近邻的索引
        nearest_indices = np.argsort(distances)[:topk]
        # 返回最近邻的索引列表
        return query_image_name, nearest_indices.tolist()
    else:
        return query_image_name, []


if __name__ == '__main__':

    data_base =  r'E:\datasets\satgeoloc_dataset'
    query_path = os.path.join(data_base, 'query')
    reference_path = os.path.join(data_base, 'reference')
    # 加载所有参考图像的路径和坐标信息
    reference_images = load_reference_images(reference_path)
    query_list = os.listdir(query_path)
    # 不使用多线程，直接迭代处理查询图像
    results = []
    for img_name in tqdm(query_list):
        result = process_query_image_center(img_name, reference_images)
        results.append(result)

    print("Saving...")
    with open(os.path.join(data_base, f"gps_dict_satgeoloc_{len(results)}querys.pkl"), "wb") as f:
        pickle.dump(results, f)
